package com.atguigu.tree;


public class RBTree<K extends Comparable<K>,V>{

    private static final boolean RED=true;
    private static final boolean BLACK=false;

    private class Node{
        public K key;
        public V value;
        public Node right;
        public Node left;
        public int height;
        public boolean color;
        public Node(K key,V value,Node left,Node right){
            this.key=key;
            this.value=value;
            this.right=right;
            this.left=left;
            this.height=1;
            this.color=RED;
        }
        public Node(K key,V value){
            this(key,value,null,null);
        }

        @Override
        public String toString() {
            return "{"+key.toString()+" : "+value.toString()+"}";
        }
    }
    private Node root;
    private int size;
    public RBTree(){
        root=null;
        size=0;
    }
    //获得node节点的高度
    private int getHeight(Node node){
        if(node==null){
            return 0;
        }
        return node.height;
    }
    //判断节点的颜色
    public boolean isRed(Node node){
        if(node==null){
            return BLACK;
        }
        return node.color;
    }
    //计算平衡因子
    private int getBalanceFactor(Node node){
        return getHeight(node.left)-getHeight(node.right);
    }
    //红黑树的左旋转
    private Node leftRotate(Node node){
        Node x=node.right;
        node.right=x.left;
        x.left=node;
        x.color=node.color;
        node.color=RED;
        return x;
    }
    //红黑树的右旋转
    private Node rightRotate(Node  node){
        Node x=node.left;
        node.left=x.right;
        x.right=node;
        x.color=node.color;
        node.color=RED;
        return x;
    }
    //颜色翻转
    private void flipColors(Node node){
        node.color=RED;
        node.left.color=BLACK;
        node.right.color=BLACK;
    }

    //向红黑树中添加新的元素
    public void add(K key, V value) {
        root=add(root,key,value);
        root.color=BLACK;//最终根节点为黑色节点
    }
    //向以root为根的红黑树中插入元素(key,value)
    //返回插入新节点后红黑树的根
    private Node add(Node root,K key,V value){
        if(root==null){
            size++;
            return new Node(key,value);
        }
        if(key.compareTo(root.key)<0){
            root.left=add(root.left,key,value);
        }else if(key.compareTo(root.key)>0){
            root.right=add(root.right,key,value);
        }else{//key.compareTo(root.key)==0
            root.value=value;
        }
        //左旋转
        if(isRed(root.right)&&!isRed(root.left)){
            root=leftRotate(root);
        }
        //右旋转
        if(isRed(root.left)&&isRed(root.left.left)){
            root=rightRotate(root);
        }
        //是否需要颜色翻转
        if(isRed(root.left)&&isRed(root.right)){
            flipColors(root);
        }

        return root;
    }


    //返回以node为根节点的二分搜索树中，key所在的节点
    private Node getNode(Node root,K key){
        if(root==null){
            return null;
        }
        if(root.key.compareTo(key)>0){
            return getNode(root.left,key);
        }else if(root.key.compareTo(key)<0){
            return getNode(root.right,key);
        }else{
            return root;
        }
    }
    private Node minimum(Node root){
        if(root.left==null){
            return root;
        }
        return minimum(root.left);
    }
    private Node removeMin(Node root){
        if(root.left==null){
            Node rightNode=root.right;
            root.right=null;
            size--;
            return rightNode;
        }
        root.left=removeMin(root.left);
        return root;
    }

    //从二分搜索树中删除Key节点
    public V remove(K key) {
        Node node=getNode(root,key);
        if(node!=null){
            root=remove(root,key);
            return node.value;
        }
        return null;
    }
    private Node remove(Node root,K key){
        if(root==null){
            return null;
        }
        Node retNode;
        if(key.compareTo(root.key)<0){
            root.left=remove(root.left,key);
            retNode=root;
        }else if(key.compareTo(root.key)>0){
            root.right=remove(root.right,key);
            retNode=root;
        }else{
            if(root.left==null){
                Node rightNode=root.right;
                size--;
                root.right=null;
                retNode=rightNode;
            }else if(root.right==null){
                Node leftNode=root.left;
                root.left=null;
                size--;
                retNode =leftNode;
            }else {
                Node successor = minimum(root.right);
                successor.right = remove(root.right,successor.key);
                successor.left = root.left;
                root.left = root.right = null;  //help GC
                retNode = successor;
            }
        }
        //若删除的是叶子节点，则retNode=null;
        if(retNode==null){
            return null;
        }
        //更新height
        retNode.height=1+Math.max(getHeight(retNode.left),getHeight(retNode.right));

        //计算平衡因子
        int balanceFactor = getBalanceFactor(retNode);
        //RR
        if(balanceFactor>1&&getBalanceFactor(retNode.left)>=0){
            return rightRotate(retNode);
        }
        //LL
        if(balanceFactor<-1&&getBalanceFactor(retNode.right)<=0){
            return leftRotate(retNode);
        }
        //RL
        if(balanceFactor<-1&&getBalanceFactor(retNode.right)>0){
            retNode.right=rightRotate(retNode.right);
            return leftRotate(retNode);
        }
        //LR
        if(balanceFactor>1&&getBalanceFactor(retNode.left)<0){
            retNode.left=leftRotate(retNode.left);
            return rightRotate(retNode);
        }
        return retNode;
    }


    public boolean contains(K key) {
        return getNode(root,key)!=null;
    }


    public V get(K key) {
        Node cur=getNode(root,key);
        return cur==null?null:cur.value;
    }


    public void set(K key, V newValue) {
        Node cur=getNode(root,key);
        if(cur!=null){
            cur.value=newValue;
        }else{
            throw new IllegalArgumentException(key+" does't exist!");
        }
    }

    public int getSize() {
        return size;
    }


    public boolean isEmpty() {
        return size==0;
    }
}
